package com.aiwiown.snackmq.storage.service;

import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.common.util.MessageConverter;
import com.aiwiown.snackmq.storage.config.StorageConfig;
import com.aiwiown.snackmq.storage.file.MappedFile;
import com.aiwiown.snackmq.storage.file.MappedFileQueue;
import com.aiwiown.snackmq.storage.index.IndexManager;
import com.aiwiown.snackmq.storage.util.ResourceManager;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.nio.Buffer;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 批量存储服务 - 提供高效的批量读写操作
 * 使用异步处理和内存池优化性能
 */
@Slf4j
public class BatchStorageService {
    
    private final StorageConfig config;
    private final MappedFileQueue mappedFileQueue;
    private final IndexManager indexManager;
    private final ExecutorService executorService;
    private final LinkedBlockingQueue<BatchWriteTask> writeQueue;
    private final AtomicLong totalMessagesWritten = new AtomicLong(0);
    private final AtomicLong totalMessagesRead = new AtomicLong(0);
    
    // 批量写入配置
    private static final int DEFAULT_BATCH_SIZE = 100;
    private static final int DEFAULT_BATCH_TIMEOUT_MS = 100;
    private static final int MAX_BATCH_SIZE = 1000;
    
    public BatchStorageService(StorageConfig config) {
        this.config = config;
        this.mappedFileQueue = new MappedFileQueue(config);
        this.indexManager = new IndexManager(config);
        this.executorService = Executors.newFixedThreadPool(
            Runtime.getRuntime().availableProcessors(),
            r -> new Thread(r, "batch-storage-worker")
        );
        this.writeQueue = new LinkedBlockingQueue<>();
        
        try {
            // 初始化组件
            this.mappedFileQueue.init();
            this.indexManager.init();
            
            // 启动批量写入处理器
            startBatchWriter();
        } catch (Exception e) {
            log.error("Failed to initialize BatchStorageService", e);
            throw new RuntimeException("Failed to initialize BatchStorageService", e);
        }
    }
    
    /**
     * 批量写入消息
     * 
     * @param messages 消息列表
     * @return 每条消息的逻辑偏移量列表
     */
    public CompletableFuture<List<Long>> writeBatch(List<Message> messages) {
        CompletableFuture<List<Long>> future = new CompletableFuture<>();
        
        try {
            // 将消息转换为字节数组
            List<byte[]> messageBytes = new ArrayList<>();
            for (Message message : messages) {
                byte[] bytes = serializeMessage(message);
                if (bytes != null) {
                    messageBytes.add(bytes);
                }
            }
            
            // 提交批量写入任务
            BatchWriteTask task = new BatchWriteTask(messageBytes, future);
            writeQueue.offer(task);
            
        } catch (Exception e) {
            log.error("Failed to prepare batch write", e);
            future.completeExceptionally(e);
        }
        
        return future;
    }
    
    /**
     * 批量读取消息
     * 
     * @param startOffset 起始逻辑偏移量
     * @param maxMessages 最大消息数量
     * @return 消息列表
     */
    public CompletableFuture<List<Message>> readBatch(long startOffset, int maxMessages) {
        return CompletableFuture.supplyAsync(() -> {
            try {
                List<Message> messages = new ArrayList<>();
                
                // 根据逻辑偏移量查找物理位置
                IndexManager.PhysicalPosition physicalPos = indexManager.findPhysicalPosition(startOffset);
                if (physicalPos == null) {
                    log.warn("Invalid start offset: {}", startOffset);
                    return messages;
                }
                
                // 获取对应的映射文件
                MappedFile targetFile = mappedFileQueue.getMappedFileByIndex(physicalPos.getFileIndex());
                if (targetFile == null) {
                    log.warn("Target file not found for fileIndex: {}", physicalPos.getFileIndex());
                    return messages;
                }
                
                log.debug("Reading from fileIndex={}, fileOffset={}, maxMessages={}", 
                        physicalPos.getFileIndex(), physicalPos.getFileOffset(), maxMessages);
                
                // 从正确位置开始批量读取
                List<ByteBuffer> messageBuffers = targetFile.readBatchMessages(physicalPos.getFileOffset(), maxMessages);
                
                log.debug("Read {} message buffers from file", messageBuffers.size());
                
                for (ByteBuffer buffer : messageBuffers) {
                    Message message = deserializeMessage(buffer);
                    if (message != null) {
                        messages.add(message);
                    }
                }
                
                totalMessagesRead.addAndGet(messages.size());
                log.debug("Successfully deserialized {} messages from offset {}", messages.size(), startOffset);
                
                return messages;
            } catch (Exception e) {
                log.error("Failed to read batch messages from offset: {}", startOffset, e);
                throw new RuntimeException(e);
            }
        }, executorService);
    }
    
    /**
     * 异步批量写入 - 使用队列缓冲
     * 
     * @param messages 消息列表
     */
    public void writeBatchAsync(List<Message> messages) {
        try {
            List<byte[]> messageBytes = new ArrayList<>();
            for (Message message : messages) {
                byte[] bytes = serializeMessage(message);
                if (bytes != null) {
                    messageBytes.add(bytes);
                }
            }
            
            BatchWriteTask task = new BatchWriteTask(messageBytes, null);
            writeQueue.offer(task);
            
        } catch (Exception e) {
            log.error("Failed to prepare async batch write", e);
        }
    }
    
    /**
     * 启动批量写入处理器
     */
    private void startBatchWriter() {
        Thread writerThread = new Thread(() -> {
            List<BatchWriteTask> batch = new ArrayList<>(MAX_BATCH_SIZE);
            while (!Thread.currentThread().isInterrupted()) {
                try {
                    // 1. 阻塞等待第一个任务，避免CPU空转
                    BatchWriteTask firstTask = writeQueue.take();
                    batch.add(firstTask);
                    // 2. 非阻塞地将队列中所有剩余任务快速取出，填充到批次中
                    writeQueue.drainTo(batch, MAX_BATCH_SIZE - 1);
                    // 3. 执行批量写入
                    if (!batch.isEmpty()) {
                        executeBatchWrite(batch);
                        batch.clear();
                    }
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    log.info("Batch writer thread interrupted.");
                    break;
                } catch (Exception e) {
                    log.error("Error in batch writer", e);
                    // 清理批次，防止任务重复处理或丢失
                    if (!batch.isEmpty()) {
                        batch.forEach(task -> {
                            if (task.getFuture() != null) {
                                task.getFuture().completeExceptionally(e);
                            }
                        });
                        batch.clear();
                    }
                }
            }
        }, "batch-writer");
        writerThread.setDaemon(true);
        writerThread.start();
    }
    
    /**
     * 执行批量写入
     */
    private void executeBatchWrite(List<BatchWriteTask> tasks) {
        try {
            // 合并所有消息字节
            List<byte[]> allMessages = new ArrayList<>();
            for (BatchWriteTask task : tasks) {
                allMessages.addAll(task.getMessageBytes());
            }
            
            // 执行批量写入
            MappedFile currentFile = mappedFileQueue.getCurrentMappedFile();
            if (currentFile != null) {
                MappedFile.BatchWriteResult result = currentFile.writeBatchMessages(allMessages);
                
                if (result.isSuccess()) {
                    int totalWritten = result.getTotalWritten();
                    totalMessagesWritten.addAndGet(totalWritten);
                    
                    // 为每条消息创建索引
                    for (MappedFile.MessageWriteInfo messageInfo : result.getMessageInfos()) {
                        try {
                            // 创建索引条目：逻辑偏移量 -> 物理位置
                            indexManager.append(messageInfo.getLogicalOffset(), 
                                             messageInfo.getFileIndex(), 
                                             messageInfo.getFileOffset());
                            
                            log.debug("Created index for message: logicalOffset={}, fileIndex={}, fileOffset={}", 
                                    messageInfo.getLogicalOffset(), messageInfo.getFileIndex(), messageInfo.getFileOffset());
                        } catch (Exception e) {
                            log.error("Failed to create index for message: {}", messageInfo, e);
                        }
                    }
                    
                    // 完成所有任务，每个任务返回自己的逻辑偏移量列表
                    for (BatchWriteTask task : tasks) {
                        if (task.getFuture() != null) {
                            // 提取当前任务对应的逻辑偏移量
                            List<Long> logicalOffsets = new ArrayList<>();
                            int messageCount = task.getMessageBytes().size();
                            int startIndex = 0;
                            
                            // 计算当前任务在总批次中的起始位置
                            for (BatchWriteTask t : tasks) {
                                if (t == task) {
                                    break;
                                }
                                startIndex += t.getMessageBytes().size();
                            }
                            
                            // 提取对应的逻辑偏移量
                            for (int i = 0; i < messageCount; i++) {
                                if (startIndex + i < result.getMessageInfos().size()) {
                                    logicalOffsets.add(result.getMessageInfos().get(startIndex + i).getLogicalOffset());
                                }
                            }
                            
                            task.getFuture().complete(logicalOffsets);
                        }
                    }
                    
                    log.debug("Batch write completed: {} messages with indices", totalWritten);
                } else {
                    // 写入失败，完成所有任务为异常
                    for (BatchWriteTask task : tasks) {
                        if (task.getFuture() != null) {
                            task.getFuture().completeExceptionally(
                                new IOException("Failed to write batch messages")
                            );
                        }
                    }
                }
            } else {
                // 没有可用的映射文件
                for (BatchWriteTask task : tasks) {
                    if (task.getFuture() != null) {
                        task.getFuture().completeExceptionally(
                            new IOException("No available mapped file for writing")
                        );
                    }
                }
            }
        } catch (Exception e) {
            log.error("Error executing batch write", e);
            for (BatchWriteTask task : tasks) {
                if (task.getFuture() != null) {
                    task.getFuture().completeExceptionally(e);
                }
            }
        }
    }
    
    /**
     * 序列化消息
     */
    private byte[] serializeMessage(Message message) {
        try {
            // 使用MessageConverter进行序列化
            return MessageConverter.toBytes(message);
        } catch (Exception e) {
            log.error("Failed to serialize message", e);
            return null;
        }
    }
    
    /**
     * 零拷贝反序列化消息
     */
    private Message deserializeMessage(ByteBuffer buffer) {
        try {
            // 检查缓冲区大小
            if (buffer.remaining() < 4) {
                log.error("Buffer too small to contain message length prefix");
                return null;
            }
            
            // 读取长度前缀
            int messageLength = buffer.getInt();
            
            // 验证消息长度
            if (messageLength <= 0) {
                log.error("Invalid message length: {}", messageLength);
                return null;
            }
            
            // 检查是否有足够的数据
            if (messageLength > buffer.remaining()) {
                log.error("Message length {} exceeds remaining buffer size {}", messageLength, buffer.remaining());
                return null;
            }
            
            // 创建消息内容的 ByteBuffer 切片（零拷贝）
            ByteBuffer messageBuffer = buffer.slice();
            // 【最终修复】: 强制类型转换为 Buffer 以兼容 JDK 8
            // 这是导致 NoSuchMethodError 的原因
            ((Buffer) messageBuffer).limit(messageLength);
            
            // 使用零拷贝反序列化
            return MessageConverter.fromByteBuffer(messageBuffer);
        } catch (Exception e) {
            log.error("Failed to deserialize message", e);
            return null;
        }
    }
    
    /**
     * 获取当前写入位置（逻辑偏移量）
     * 
     * @return 当前写入位置
     */
    public long getCurrentWritePosition() {
        try {
            MappedFile currentFile = mappedFileQueue.getCurrentMappedFile();
            if (currentFile == null) {
                return 0;
            }
            
            // 计算当前文件的逻辑偏移量
            int fileIndex = mappedFileQueue.getMappedFiles().indexOf(currentFile);
            if (fileIndex == -1) {
                return 0;
            }
            
            long mappedFileSize = config.getMappedFileSize();
            return (long) fileIndex * mappedFileSize + currentFile.getWrotePosition();
        } catch (Exception e) {
            log.error("Error getting current write position", e);
            return 0;
        }
    }
    
    /**
     * 获取统计信息
     */
    public BatchStorageStats getStats() {
        return new BatchStorageStats(
            totalMessagesWritten.get(),
            totalMessagesRead.get(),
            writeQueue.size()
        );
    }
    
    /**
     * 关闭服务
     */
    public void close() {
        try {
            // 等待所有正在进行的批量写入操作完成
            log.info("Waiting for pending batch write operations to complete...");
            int maxWaitTime = 10; // 最多等待10秒
            int waitCount = 0;
            while (!writeQueue.isEmpty() && waitCount < maxWaitTime) {
                try {
                    Thread.sleep(1000);
                    waitCount++;
                    log.debug("Waiting for batch operations to complete, queue size: {}", writeQueue.size());
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    break;
                }
            }
            
            if (!writeQueue.isEmpty()) {
                log.warn("Some batch operations may not have completed, remaining: {}", writeQueue.size());
            }
            
            // 使用资源管理器安全关闭线程池
            ResourceManager.safeShutdownExecutor(executorService, 5);
            
            // 关闭文件队列和索引管理器
            if (mappedFileQueue != null) {
                mappedFileQueue.close();
            }
            if (indexManager != null) {
                indexManager.close();
            }
            
            // 强制垃圾回收以确保MappedByteBuffer被释放
            ResourceManager.forceGcAndWait(1000);
            
        } catch (Exception e) {
            log.error("Error closing batch storage service", e);
        }
    }
    
    /**
     * 批量写入任务
     */
    private static class BatchWriteTask {
        private final List<byte[]> messageBytes;
        private final CompletableFuture<List<Long>> future;
        
        public BatchWriteTask(List<byte[]> messageBytes, CompletableFuture<List<Long>> future) {
            this.messageBytes = messageBytes;
            this.future = future;
        }
        
        public List<byte[]> getMessageBytes() {
            return messageBytes;
        }
        
        public CompletableFuture<List<Long>> getFuture() {
            return future;
        }
    }
    
    /**
     * 批量存储统计信息
     */
    public static class BatchStorageStats {
        private final long totalMessagesWritten;
        private final long totalMessagesRead;
        private final int pendingWrites;
        
        public BatchStorageStats(long totalMessagesWritten, long totalMessagesRead, int pendingWrites) {
            this.totalMessagesWritten = totalMessagesWritten;
            this.totalMessagesRead = totalMessagesRead;
            this.pendingWrites = pendingWrites;
        }
        
        public long getTotalMessagesWritten() {
            return totalMessagesWritten;
        }
        
        public long getTotalMessagesRead() {
            return totalMessagesRead;
        }
        
        public int getPendingWrites() {
            return pendingWrites;
        }
    }
} 